from core.cuda import cuda_module
from core.optimizer import Optimizer


class AdaDelta(Optimizer):
    def __init__(self, rho=0.95, eps=1e-6):
        """
        AdaDelta 优化器，用于自适应调整学习率。

        Args:
            rho (float): 平滑系数，一般取 0.9 到 0.999。
            eps (float): 为了数值稳定性而添加的小常数，一般取 1e-6。
        """
        super().__init__()
        self.rho = rho
        self.eps = eps
        self.msg = {}   # 存储梯度的平方的滑动平均
        self.msdx = {}  # 存储参数更新的平方的滑动平均

    def update_one(self, param):
        """
        更新单个参数。

        Args:
            param (Parameter): 待更新的参数。

        Returns:
            None
        """
        xp = cuda_module

        key = id(param)
        if key not in self.msg:
            self.msg[key] = xp.zeros_like(param.data)
            self.msdx[key] = xp.zeros_like(param.data)

        msg, msdx = self.msg[key], self.msdx[key]
        rho = self.rho
        eps = self.eps
        grad = param.grad.data

        msg *= rho
        msg += (1 - rho) * grad * grad
        dx = xp.sqrt((msdx + eps) / (msg + eps)) * grad
        msdx *= rho
        msdx += (1 - rho) * dx * dx
        param.data -= dx

